from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

import torchvision
from torchvision import transforms

"""
使用torchvision的数据集 并使用 Dataloader加载数据集合
"""
# 加载的数据集 图片格式是PIL.Image  转成成 Tensor
transform  = transforms.Compose([
    transforms.ToTensor()
])

# tensorboard 面板
writer = SummaryWriter("logs04")

# 训练集合 和 测试集合
trainSet = torchvision.datasets.CIFAR10("../dataset/CIFAR10" , transform=transform ,train=True , download=True)
#testSet = torchvision.datasets.CIFAR10("../dataset/CIFAR10" , transform=transform ,train=False , download=True)

#测试集第一涨图片
img , target = trainSet[0]
print(img.shape)
print(target)
print(trainSet.classes[target])

"""
dataset     数据集合
batch_size  批处理大小
shuffle     是否洗牌  每次取不按照顺序
drop_last   是否舍弃最后不足batch_size大小的数据集
num_workers 处理数据的线程数
"""
dataloader = DataLoader(dataset=trainSet , batch_size= 64 , shuffle= True , drop_last=False ,num_workers=0)

print(dataloader)
step = 1
for data in dataloader:
    imgs , target = data
    writer.add_images("images" , imgs , global_step=step)
    step+=1

# 关闭流
writer.close()


